from pytorch_lightning.loggers import CSVLogger
import inspect


class_registry = {
    'CSVLogger':CSVLogger,
}

def register_class(name=None, ):
    """
    A decorator to register a class in the class_registry with optional additional keys.

    Args:
        name (str or list of str, optional): The name(s) to be used as the key(s) in the registry.
                                             Defaults to cls.__name__ if not provided.

    Returns:
        function: The decorator function that registers the class.
    """
    def decorator(cls):
        if name is None:
            names = [cls.__name__]
        elif isinstance(name, str):
            names = [cls.__name__, name]
        else:
            names = name + [cls.__name__]

        for key in names:
            class_registry[key] = cls

        return cls

    return decorator


def class_builder(config, **kw):

    assert 'type' in config, "The configuration must contain a 'type' key."
    assert 'args' in config, "The configuration must contain an 'args' key."
    # print(class_registry.keys())
    assert config['type'] in class_registry, f"class {config['type']} not registered in class_registry ({class_registry.keys()})"

    CLASS_NAME = class_registry[config['type']]
    # Get the signature of the __init__ method
    init_signature = inspect.signature(CLASS_NAME.__init__)

    for k in ['normalizer_y', 'normalizer_x' , 'save_dir']:
        if k in init_signature.parameters:
            if k in kw:
                config['args'][k] = kw[k]
            else:
                print(f"{k} is not provided by class_builder, but {config['type']} requires it.")
            
    return CLASS_NAME(**config['args'])

